package algorithm.leetcode.medium;

/*
 实现一个 Trie (前缀树)，包含 insert, search, 和 startsWith 这三个操作。
 */

import java.util.TreeMap;

/**
 * @author jack.wu
 * @version 1.0
 * @date 2020/12/16
 */
public class Question208 {

    static class Trie {
        private Node root;

        private class Node {
            private TreeMap<Character, Node> next;
            private boolean isWord;

            public Node(boolean isWord) {
                next = new TreeMap<>();
                this.isWord = isWord;
            }

            public Node() {
                this(false);
            }
        }


        public Trie() {
            root = new Node();
        }

        public void insert(String word) {
            Node cur = this.root;
            char ch;
            for (int i = 0; i < word.length(); i++) {
                ch = word.charAt(i);
                cur.next.putIfAbsent(ch, new Node());
                cur = cur.next.get(ch);
            }
            if (!cur.isWord) {
                cur.isWord = true;
            }
        }

        public boolean search(String word) {
            Node cur = root;
            char ch;
            for (int i = 0; i < word.length(); i++) {
                ch = word.charAt(i);
                if (!cur.next.containsKey(ch)) {
                    return false;
                }
                cur = cur.next.get(ch);
            }
            return cur.isWord;
        }


        public boolean startsWith(String prefix) {
            Node cur = this.root;
            char ch;
            for (int i = 0; i < prefix.length(); i++) {
                ch = prefix.charAt(i);
                if (!cur.next.containsKey(ch)) {
                    return false;
                }
                cur = cur.next.get(ch);
            }
            return true;
        }
    }

    public static void main(String[] args) {

        Trie trie = new Trie();
        trie.insert("apple");
        System.out.println(trie.search("apple"));
        System.out.println(trie.search("app"));
        System.out.println(trie.startsWith("app"));
        trie.insert("app");
        System.out.println( trie.search("app"));
    }
}
